#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
处理url文件为对象，支持图像、音频和PDF文件
通过字典设置文件类型
"""

import os
import requests
import tempfile
import shutil
import uuid
import io
from pathlib import Path
from urllib.parse import urlparse


class FileURLHandler:
    """原理：下载文件后，加载为对象；"""
    # 文件类型定义
    FILE_TYPES = {
        'IMAGE': 'image',
        'AUDIO': 'audio',
        'PDF': 'pdf'
    }
    
    def __init__(self, temp_dir="./temp"):
        # 创建临时目录
        self.temp_dir = Path(temp_dir)
        if not self.temp_dir.exists():
            self.temp_dir.mkdir(parents=True)
    
    def download_file(self, url):
        """从URL下载文件到临时目录"""
        try:
            # 获取文件名，如果URL中没有文件名，则生成一个随机文件名
            parsed_url = urlparse(url)
            filename = os.path.basename(parsed_url.path)
            if not filename:
                filename = str(uuid.uuid4())
            
            # 发送请求获取文件
            response = requests.get(url, stream=True)
            if response.status_code != 200:
                return None, f"下载失败，状态码: {response.status_code}"
            
            # 保存到临时目录
            temp_file_path = self.temp_dir / filename
            with open(temp_file_path, 'wb') as f:
                for chunk in response.iter_content(chunk_size=8192):
                    f.write(chunk)
            
            return str(temp_file_path), None
        except Exception as e:
            return None, f"下载文件时出错: {str(e)}"
    
    def process_image(self, file_path):
        """处理图像文件"""
        try:
            from PIL import Image
            # 打开图像并将其加载到内存
            with Image.open(file_path) as img:
                # 创建图像的副本，避免依赖原始文件
                img_copy = img.copy()
            return img_copy, None
        except ImportError:
            return None, "请安装PIL库: pip install Pillow"
        except Exception as e:
            return None, f"处理图像时出错: {str(e)}"
    
    def process_audio(self, file_path):
        """处理音频文件"""
        try:
            import librosa
            import numpy as np
            
            # librosa.load 会将整个音频文件加载到内存中，所以删除文件不会有影响
            audio, sr = librosa.load(file_path, sr=None)
            
            # 为确保安全，可以创建数据的副本，但通常不需要
            audio_copy = np.copy(audio)
            return (audio_copy, sr), None
        except ImportError:
            return None, "请安装librosa库: pip install librosa"
        except Exception as e:
            return None, f"处理音频时出错: {str(e)}"
    
    def process_pdf(self, file_path):
        """处理PDF文件"""
        try:
            import PyPDF2
            
            # 读取文件内容到内存缓冲区
            with open(file_path, 'rb') as file:
                pdf_content = io.BytesIO(file.read())
            
            # 从内存缓冲区创建PdfReader对象
            pdf_reader = PyPDF2.PdfReader(pdf_content)
            
            # 创建一个简单封装以跟踪内存对象，防止垃圾回收
            class PdfObject:
                def __init__(self, reader, buffer):
                    self.reader = reader
                    self.buffer = buffer  # 保持对buffer的引用
                    
                def __getattr__(self, name):
                    # 代理所有属性访问到reader
                    return getattr(self.reader, name)
            
            return PdfObject(pdf_reader, pdf_content), None
        except ImportError:
            return None, "请安装PyPDF2库: pip install PyPDF2"
        except Exception as e:
            return None, f"处理PDF时出错: {str(e)}"
    
    def cleanup(self, file_path):
        """删除临时文件"""
        try:
            if os.path.exists(file_path):
                os.remove(file_path)
            return True
        except Exception as e:
            print(f"删除临时文件时出错: {str(e)}")
            return False
    
    def process_file_url(self, url, file_type=None, auto_cleanup=True):
        """
        处理文件URL
        :param url: 文件URL
        :param file_type: 文件类型，使用 FILE_TYPES 字典中的值 ('image', 'audio', 'pdf')
        :param auto_cleanup: 处理完成后是否自动删除临时文件
        :return: (处理后的对象, 临时文件路径, 错误信息)
        """
        # 默认为图像类型
        if file_type is None:
            file_type = self.FILE_TYPES['IMAGE']
        
        # 下载文件
        file_path, error = self.download_file(url)
        if error:
            return None, None, error
        
        # 根据文件类型处理
        obj, error = None, None
        try:
            if file_type == self.FILE_TYPES['IMAGE']:
                obj, error = self.process_image(file_path)
            elif file_type == self.FILE_TYPES['AUDIO']:
                obj, error = self.process_audio(file_path)
            elif file_type == self.FILE_TYPES['PDF']:
                obj, error = self.process_pdf(file_path)
            else:
                error = f"不支持的文件类型: {file_type}"
            
            # 只有在成功处理文件后，如果设置了自动清理，才删除临时文件
            if not error and auto_cleanup:
                self.cleanup(file_path)
                if auto_cleanup:
                    file_path = None
                    
        except Exception as e:
            error = f"处理文件时出错: {str(e)}"
            # 如果处理失败，清理文件
            self.cleanup(file_path)
            file_path = None
            
        return obj, file_path, error


# 使用示例
if __name__ == "__main__":
    handler = FileURLHandler()
    
    # 示例：处理音频文件
    wav_file = "https://work8.top/static/uploads/2025/06/56988dd025da45e3828c69f73d7abb98.flac"
    audio_obj, temp_path, error = handler.process_file_url(wav_file, file_type=handler.FILE_TYPES['AUDIO'], auto_cleanup=True)
    
    if error:
        print(f"处理音频文件出错: {error}")
    else:
        print(f"音频文件处理成功")
        print(f"音频对象: {audio_obj}")
        
        # 测试：确认对象是否可用（即使临时文件已删除）
        if audio_obj is not None:
            audio_data, sample_rate = audio_obj
            print(f"音频长度: {len(audio_data)}, 采样率: {sample_rate}")
            
    # 示例：处理PDF文件
    pdf_file = "https://work8.top/static/uploads/2025/06/c85af7dbfa13400cb15547611f181b93.pdf"
    print(f"\n尝试处理PDF文件: {pdf_file}")
    pdf_obj, temp_path, error = handler.process_file_url(pdf_file, file_type=handler.FILE_TYPES['PDF'], auto_cleanup=True)
    
    if error:
        print(f"处理PDF文件出错: {error}")
    else:
        print(f"PDF文件处理成功")
        print(f"PDF页数: {len(pdf_obj.reader.pages)}")
            

